%%% Copyright 2010 Google Inc.
%%%
%%% Licensed under the Apache License, Version 2.0 (the "License");
%%% you may not use this file except in compliance with the License.
%%% You may obtain a copy of the License at
%%%
%%%     http://www.apache.org/licenses/LICENSE-2.0
%%%
%%% Unless required by applicable law or agreed to in writing, software
%%% distributed under the License is distributed on an "AS IS" BASIS,
%%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%%% See the License for the specific language governing permissions and
%%% limitations under the License.

%%% @copyright 2010 Google Inc. All Rights Reserved.
%%% @author <jacoby@google.com> Jacoby Thwaites [http://sites/madcloud]
%%% @version {@version}
%%% @doc
%%%
%%% Functions on sets of sets of name-ref pairs (nr_sets). This is used
%%% to deal with output sets of operations where each set has a name
%%% and we are interested in the union of their output sets.
%%%
%%%
-module(nr_set_dict).

-include("txm.hrl").

-export([new/0, from_proplist/1, from_dict/1, get_nr_set/2, store_nr_set/3]).
-export([ref_union_set/1, get_nr_set_names/1]).
-export([count_sets/1, count_refs/1]).
-export([fold/3]).

%% @doc Returns a new, empty nr_set_dict.
%% @spec new() -> nr_set_dict()
-spec(new() -> nr_set_dict()).
new() -> #?NR_SET_DICT{}.

%% @doc Returns a new nr_set_dict populated using the proplist-of-proplists
%% which delegates to nr_set:from_proplist/1 for each set.
%% @spec from_proplist(list()) -> nr_set_dict()
-spec(from_proplist(list({string(), any()})) -> nr_set_dict()).
from_proplist(Proplist) ->
  lists:foldl(fun({Name, SetProplist}, Acc) ->
                  NrSet = nr_set:from_proplist(SetProplist),
                  store_nr_set(Name, NrSet, Acc)
             end, new(), Proplist).

%% @doc Returns a new nr_set_dict populated using `DictOfDicts'
%% which delegates to nr_set:from_dict_key_name/1 for each dict.
%% @spec from_dict(dict()) -> nr_set_dict()
-spec(from_dict(dict()) -> nr_set_dict()).
from_dict(DictOfDicts) ->
  dict:fold(fun(Name, SetDict, Acc) ->
                 NrSet = nr_set:from_dict_key_name(SetDict),
                 store_nr_set(Name, NrSet, Acc)
            end, new(), DictOfDicts).

%% @doc Stores the supplied nr_set into the nr_set dict, replacing any
%% entry already present under that name.
%% @spec store_nr_set(string(), nr_set(), nr_set_dict()) -> nr_set_dict()
-spec(store_nr_set(string(), nr_set(), nr_set_dict()) -> nr_set_dict()).
store_nr_set(Name, NrSet, #?NR_SET_DICT{
                           dict = Dict,
                           refSet = RefSet} = NrSetDict) ->
  NewDict = dict:store(Name, NrSet, Dict),
  NrSetDict#?NR_SET_DICT{
            dict = NewDict,
            refSet =
              % If there's an old entry under this name, re-create the
              % union set to ensure we're not leaving old references in it.
              case dict:find(Name, Dict) of
                error ->
                  sets:union(RefSet, nr_set:ref_set(NrSet));
                {ok, _OldEntry} ->
                  dict:fold(fun(_Name, NrSet1, Acc) ->
                                 sets:union(nr_set:ref_set(NrSet1), Acc)
                            end, sets:new(), NewDict)
                end}.

%% @doc Gets the nr_set stored under the supplied name, or not_found if
%% not in the dict.
%% @spec get_nr_set(string(), nr_set_dict()) -> nr_set()
-spec(get_nr_set(string(), nr_set_dict()) -> nr_set()).
get_nr_set(Name, #?NR_SET_DICT{
                  dict = Dict}) ->
  case dict:find(Name, Dict) of
    {ok, NrSet} ->
      NrSet;
    error -> ?NOT_FOUND
    end.

%% @doc Returns the list of nr_set names stored in the nr_set_dict.
%% @spec get_nr_set_names(nr_set_dict()) -> list(string())
-spec(get_nr_set_names(nr_set_dict()) -> list(string)).
get_nr_set_names(#?NR_SET_DICT{
                  dict = Dict}) ->
  dict:fetch_keys(Dict).

%% @doc Returns the union set of all refs in the stored ns_sets.
%% @spec ref_union_set(nr_set_dict()) -> set()
-spec(ref_union_set(nr_set_dict()) -> set()).
ref_union_set(#?NR_SET_DICT{
                   refSet = RefSet}) ->
  RefSet.

%% @doc Returns the count of the nr_set entries in this nr_set_dict.
%% @spec count_sets(nr_set_dict()) -> integer()
-spec(count_sets(nr_set_dict()) -> integer()).
count_sets(#?NR_SET_DICT{dict = Dict}) ->
  dict:size(Dict).

%% @doc Returns the count of the union of field refs in this nr_set_dict.
%% @spec count_refs(nr_set_dict()) -> integer()
-spec(count_refs(nr_set_dict()) -> integer()).
count_refs(#?NR_SET_DICT{refSet = RefSet}) ->
  sets:size(RefSet).


%% @doc Folds the supplied `Fun/3' over the dictionary of nr_sets and
%% returns the resulting accumulator.
%% The fun spec is `fun(name(), nr_set(), term()) -> term()' where
%% term() is the accumulator.
%% @spec fold(fun(), term(), nr_set_dict()) -> term()
-spec(fold(fun(), term(), nr_set_dict()) -> term()).
fold(Fun, InitialAcc, #?NR_SET_DICT{dict = Dict}) ->
  dict:fold(Fun, InitialAcc, Dict).
